#js转ast#
import os
import sys
import json
import tempfile
import subprocess
import logging
import argparse
from pathlib import Path
from typing import Dict, Any, Optional, Tuple

# 设置控制台编码为UTF-8
sys.stdout.reconfigure(encoding='utf-8')
os.system('chcp 65001 >nul')

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class SelectiveAstGenerator:
    """选择性AST生成器 - 支持单独处理benign或malicious"""
    
    def __init__(self):
        #self.extracted_dir = Path(r'D:\PythonProject\malicious_npm_detector_1.0.1\data\extracted_packages')
        #self.ast_dir = Path(r'E:\data')
        
        self.extracted_dir = Path(r'F:\robust_extracted_packages')
        self.ast_dir = Path(r'F:\robust_packages_ast')
        self.bridge_script = self._get_bridge_script_path()
        
        # 创建AST目录
        self.ast_dir.mkdir(parents=True, exist_ok=True)
        (self.ast_dir / 'benign').mkdir(exist_ok=True)
        (self.ast_dir / 'malicious').mkdir(exist_ok=True)
        
    def _get_bridge_script_path(self) -> str:
        """获取Node.js桥接脚本路径"""
        script_dir = Path(__file__).parent / 'js'
        return str(script_dir / 'babel_parser_bridge.js')
    
    def _create_bridge_script(self):
        """创建Node.js桥接脚本"""
        script_dir = Path(__file__).parent / 'js'
        script_dir.mkdir(exist_ok=True)
        
        bridge_script_content = '''
const fs = require('fs');

// 使用Babel解析器
function parseWithBabel(code) {
    const babel = require('@babel/parser');
    return babel.parse(code, {
        sourceType: 'unambiguous',
        allowImportExportEverywhere: true,
        allowReturnOutsideFunction: true,
        plugins: [
            'jsx',
            'typescript',
            'decorators-legacy',
            'classProperties',
            'objectRestSpread',
            'asyncGenerators',
            'dynamicImport',
            'exportDefaultFrom',
            'exportNamespaceFrom'
        ]
    });
}

// 主函数
function main() {
    const inputFile = process.argv[2];
    const outputFile = process.argv[3];
    
    if (!inputFile || !outputFile) {
        console.error('用法: node babel_parser_bridge.js <input.js> <output.json>');
        process.exit(1);
    }
    
    try {
        const code = fs.readFileSync(inputFile, 'utf8');
        const ast = parseWithBabel(code);
        
        const result = {
            success: true,
            ast: ast,
            parser: 'babel',
            error: null
        };
        
        fs.writeFileSync(outputFile, JSON.stringify(result, null, 2));
        
    } catch (error) {
        const result = {
            success: false,
            ast: null,
            parser: 'babel',
            error: error.message
        };
        fs.writeFileSync(outputFile, JSON.stringify(result, null, 2));
    }
}

main();
'''
        
        bridge_script_path = script_dir / 'babel_parser_bridge.js'
        with open(bridge_script_path, 'w', encoding='utf-8') as f:
            f.write(bridge_script_content)
        
        return str(bridge_script_path)
    
    def _clean_unicode_surrogates(self, obj):
        """清理Unicode代理对字符"""
        if isinstance(obj, dict):
            return {key: self._clean_unicode_surrogates(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [self._clean_unicode_surrogates(item) for item in obj]
        elif isinstance(obj, str):
            # 移除无效的Unicode代理对
            try:
                # 尝试编码为UTF-8，如果失败则清理
                obj.encode('utf-8')
                return obj
            except UnicodeEncodeError:
                # 移除所有代理对字符
                import re
                return re.sub(r'[\ud800-\udfff]', '?', obj)
        else:
            return obj
    
    def generate_ast_for_file(self, js_file_path: Path) -> Tuple[bool, Optional[Dict], str]:
        """为单个JS文件生成AST"""
        try:
            # 检查文件大小（限制10MB）
            if js_file_path.stat().st_size > 10 * 1024 * 1024:
                return False, None, "文件过大(>10MB)"
            
            # 读取JS代码
            with open(js_file_path, 'r', encoding='utf-8', errors='ignore') as f:
                js_code = f.read()
            
            # 创建临时文件
            with tempfile.NamedTemporaryFile(suffix='.js', delete=False, mode='w', encoding='utf-8') as temp_js:
                temp_js_path = temp_js.name
                temp_js.write(js_code)
            
            temp_ast_path = f"{temp_js_path}_ast.json"
            
            try:
                # 确保桥接脚本存在
                if not os.path.exists(self.bridge_script):
                    self.bridge_script = self._create_bridge_script()
                
                # 调用Node.js脚本
                result = subprocess.run(
                    ['node', self.bridge_script, temp_js_path, temp_ast_path],
                    capture_output=True,
                    text=True,
                    timeout=30
                )
                
                if result.returncode == 0 and os.path.exists(temp_ast_path):
                    with open(temp_ast_path, 'r', encoding='utf-8') as f:
                        result_data = json.load(f)
                    
                    if result_data.get('success'):
                        # 清理AST中的Unicode问题
                        clean_ast = self._clean_unicode_surrogates(result_data.get('ast'))
                        return True, clean_ast, 'babel'
                    else:
                        return False, None, result_data.get('error', '解析失败')
                else:
                    error_msg = result.stderr.strip() if result.stderr else "Node.js执行失败"
                    return False, None, error_msg
                    
            finally:
                # 清理临时文件
                for temp_file in [temp_js_path, temp_ast_path]:
                    if os.path.exists(temp_file):
                        os.unlink(temp_file)
                        
        except Exception as e:
            return False, None, str(e)
    
    def save_ast_or_error(self, ast: Optional[Dict], success: bool, error_reason: str, category: str, package_name: str, relative_js_path: Path) -> str:
        """保存AST到文件，如果失败则保存错误信息"""
        # 构建完整的目标路径，保持原有目录结构
        target_dir = self.ast_dir / category / package_name / relative_js_path.parent
        target_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成AST文件名：原文件名.ast.json
        ast_filename = f"{relative_js_path.name}.ast.json"
        ast_file_path = target_dir / ast_filename
        
        if success and ast:
            # 保存成功的AST
            ast_data = {
                "success": True,
                "parser": "babel",
                "ast": ast,
                "error": None
            }
        else:
            # 保存失败信息
            ast_data = {
                "success": False,
                "parser": "babel",
                "ast": None,
                "error": error_reason
            }
        
        # 清理所有数据中的Unicode问题
        clean_ast_data = self._clean_unicode_surrogates(ast_data)
        
        try:
            # 保存到文件，使用更安全的编码方式
            with open(ast_file_path, 'w', encoding='utf-8', errors='replace') as f:
                json.dump(clean_ast_data, f, ensure_ascii=False, indent=2)
        except UnicodeEncodeError:
            # 如果还是有问题，使用ASCII编码
            with open(ast_file_path, 'w', encoding='utf-8') as f:
                json.dump(clean_ast_data, f, ensure_ascii=True, indent=2)
        
        return str(ast_file_path)
    
    def _get_javascript_files(self, directory: Path) -> list:
        """获取目录下的所有JavaScript文件（.js, .cjs, .mjs），排除以这些扩展名结尾的目录"""
        js_files = []
        for pattern in ["*.js", "*.cjs", "*.mjs"]:
            # 使用rglob递归搜索，但只包含文件，排除目录
            js_files.extend([f for f in directory.rglob(pattern) if f.is_file()])
        return js_files
    
    def get_category_stats(self, category: str) -> Dict[str, int]:
        """获取指定类别的统计信息"""
        category_dir = self.extracted_dir / category
        
        if not category_dir.exists():
            return {'packages': 0, 'js_files': 0}
        
        packages = 0
        js_files = 0
        
        for package_dir in category_dir.iterdir():
            if package_dir.is_dir():
                packages += 1
                # 使用辅助方法获取JavaScript文件
                js_files += len(self._get_javascript_files(package_dir))
        
        return {'packages': packages, 'js_files': js_files}
        
        return {'packages': packages, 'js_files': js_files}
    
    def process_category(self, category: str):
        """处理指定类别的包"""
        if category not in ['benign', 'malicious']:
            raise ValueError(f"无效的类别: {category}。必须是 'benign' 或 'malicious'")
        
        category_dir = self.extracted_dir / category
        
        if not category_dir.exists():
            print(f"❌ 目录不存在: {category_dir}")
            return
        
        # 获取统计信息
        stats = self.get_category_stats(category)
        
        print(f"🎯 开始处理 {category.upper()} 类别")
        print(f"📊 预览统计: {stats['packages']} 个包, {stats['js_files']} 个JavaScript文件(.js/.cjs/.mjs)")
        print(f"🔧 使用解析器: @babel/parser")
        print(f"📁 保持完整目录结构...")
        print("="*60)
        
        # 创建一个专门的日志文件来记录所有失败项
        failure_log_path = self.ast_dir / f"ast_generation_failures_{category}.log"
        with open(failure_log_path, 'w', encoding='utf-8') as f:
            f.write(f"--- {category.upper()} 类别AST生成失败详细报告 ---\n\n")
        
        total_files = 0
        success_count = 0
        failed_count = 0
        failed_files = []
        processed_packages = 0
        
        # 遍历包
        for package_dir in category_dir.iterdir():
            if not package_dir.is_dir():
                continue
            
            package_name = package_dir.name
            processed_packages += 1
            
            print(f"\n📦 [{processed_packages}/{stats['packages']}] 处理包: {category}/{package_name}")
            
            # 使用辅助方法查找所有JavaScript相关文件（.js, .cjs, .mjs），但排除以这些扩展名结尾的目录
            js_files = self._get_javascript_files(package_dir)
            
            if not js_files:
                print(f"   ℹ️ 无JavaScript文件")
                continue
            
            print(f"   📄 发现 {len(js_files)} 个JavaScript文件")
            
            for i, js_file in enumerate(js_files, 1):
                total_files += 1
                
                # 生成相对路径（保持目录结构）
                relative_path = js_file.relative_to(package_dir)
                
                print(f"   [{i}/{len(js_files)}] 处理: {relative_path}")
                
                # 生成AST
                success, ast, info = self.generate_ast_for_file(js_file)
                
                # 无论成功还是失败，都保存结果
                try:
                    ast_path = self.save_ast_or_error(ast, success, info, category, package_name, relative_path)
                    
                    if success and ast:
                        success_count += 1
                        print(f"      ✅ AST已保存")
                    else:
                        failed_count += 1
                        failed_files.append({
                            'file': f"{category}/{package_name}/{relative_path}",
                            'reason': info,
                            'category': category,
                            'package': package_name,
                            'relative_path': str(relative_path)
                        })
                        print(f"      ❌ 失败: {info}")
                except Exception as e:
                    failed_count += 1
                    error_msg = f"保存失败: {str(e)}"
                    failed_files.append({
                        'file': f"{category}/{package_name}/{relative_path}",
                        'reason': error_msg,
                        'category': category,
                        'package': package_name,
                        'relative_path': str(relative_path)
                    })
                    print(f"      ❌ 保存失败: {error_msg}")
        
        # 输出详细统计信息，并传入日志文件路径
        self._print_category_summary(category, total_files, success_count, failed_count, failed_files, processed_packages, failure_log_path)
    def _print_category_summary(self, category: str, total_files: int, success_count: int, 
                                   failed_count: int, failed_files: list, processed_packages: int,
                                   failure_log_path: Path):
        """打印类别处理摘要，并将完整失败列表写入日志文件"""
        summary_lines = [
            "\n" + "="*80,
            f"🎯 {category.upper()} 类别处理完成",
            "="*80,
            f"📦 处理包数: {processed_packages}",
            f"📄 总文件数: {total_files}",
            f"✅ 成功生成: {success_count}",
            f"❌ 失败: {failed_count}",
            f"📊 成功率: {(success_count/total_files*100):.1f}%" if total_files > 0 else "📊 成功率: 0%",
            f"\nℹ️ 完整的失败文件列表已保存到: {failure_log_path}"
        ]
        
        print("\n".join(summary_lines))
        
        if failed_files:
            # 按原因分组
            reasons = {}
            for failed in failed_files:
                reason = failed['reason']
                if reason not in reasons:
                    reasons[reason] = []
                reasons[reason].append(failed)
            
            # 将完整、未省略的列表写入日志文件
            with open(failure_log_path, 'a', encoding='utf-8') as f:
                f.write("\n--- 失败原因详细分析 ---\n")
                # 按失败文件数量对原因进行排序
                for reason, files in sorted(reasons.items(), key=lambda item: len(item[1]), reverse=True):
                    f.write(f"\n❌ 原因: [{reason}] - {len(files)}个文件\n")
                    for file_info in files:
                        f.write(f"   - {file_info['file']}\n")
    
            # 在控制台打印一个简短的预览
            print(f"\n🔍 控制台失败原因预览 (完整列表见日志文件):")
            # 按失败文件数量对原因进行排序
            for reason, files in sorted(reasons.items(), key=lambda item: len(item[1]), reverse=True)[:5]: # 只预览前5大原因
                print(f"\n❌ [{reason}] - {len(files)}个文件:")
                for file_info in files[:3]: # 每个原因只预览前3个
                    print(f"   - {file_info['file']}")
                if len(files) > 3:
                    print(f"   ... (及其他 {len(files) - 3} 个文件)")

    def process_both_categories(self):
        """处理所有类别"""
        print("🚀 开始处理所有JS文件，生成AST...")
        
        for category in ['benign', 'malicious']:
            self.process_category(category)
        
        print(f"\n🎉 所有类别处理完成!")

def main():
    """主函数 - 支持命令行参数"""
    parser = argparse.ArgumentParser(description='AST生成器 - 支持选择性处理')
    parser.add_argument('--category', '-c', 
                       choices=['benign', 'malicious', 'both'], 
                       default='both',
                       help='选择要处理的类别: benign(良性), malicious(恶意), both(全部)')
    parser.add_argument('--preview', '-p', 
                       action='store_true',
                       help='预览统计信息，不执行实际处理')
    
    args = parser.parse_args()
    
    generator = SelectiveAstGenerator()
    
    # 预览模式
    if args.preview:
        print("🔍 预览模式 - 统计信息")
        print("="*50)
        
        for category in ['benign', 'malicious']:
            stats = generator.get_category_stats(category)
            print(f"{category.upper()}: {stats['packages']} 个包, {stats['js_files']} 个JS文件")
        
        return
    
    # 执行处理
    if args.category == 'both':
        generator.process_both_categories()
    else:
        generator.process_category(args.category)

if __name__ == "__main__":
    main()